#include <ntddk.h>
#include <windef.h>
#include <tchar.h>

#include "../NTModelDrvCommon/NTModelDrvCommon.h"

#define DEVICE_NAME           L"\\Device\\Ntmodeldrv"
#define SYMBOLIC_LINK_NAME    L"\\??\\Ntmodeldrv"

VOID
DriverUnload(PDRIVER_OBJECT DriverObject);

NTSTATUS
NtModelDrvCommon(PDEVICE_OBJECT DeviceObject, PIRP Irp);

NTSTATUS
NtModelDrvCreateClose(PDEVICE_OBJECT DeviceObject, PIRP Irp);

NTSTATUS
NtModelDrvRead(PDEVICE_OBJECT DeviceObject, PIRP Irp);

NTSTATUS
NtModelDrvWrite(PDEVICE_OBJECT DeviceObject, PIRP Irp);

NTSTATUS
NtModelDrvIoControl(PDEVICE_OBJECT DeviceObject, PIRP Irp);

NTSTATUS
DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING RegisterPath)
{
	NTSTATUS Status            = STATUS_SUCCESS;
	DriverObject->DriverUnload = DriverUnload;

	PDEVICE_OBJECT DeviceObject = NULL;
	do {

		UNICODE_STRING DeviceName = { 0 };
		RtlInitUnicodeString(&DeviceName, DEVICE_NAME);

		Status = IoCreateDevice(DriverObject, 0, &DeviceName, FILE_DEVICE_UNKNOWN, 0, FALSE, &DeviceObject);
		if (NT_SUCCESS(Status) == FALSE) {

			KdPrintEx((77, 0, "IoCreateDevice Failed: %x\n", Status));
			break;
		}
		DeviceObject->Flags |= DO_BUFFERED_IO;

		UNICODE_STRING SymbolicLinkName = { 0 };
		RtlInitUnicodeString(&SymbolicLinkName, SYMBOLIC_LINK_NAME);

		Status = IoCreateSymbolicLink(&SymbolicLinkName, &DeviceName);
		if (NT_SUCCESS(Status) == FALSE) {

			IoDeleteDevice(DeviceObject);
			KdPrintEx((77, 0, "IoCreateSymbolicLink Failed: %x\n", Status));
			break;
		}

		for (int i = 0; i < IRP_MJ_MAXIMUM_FUNCTION + 1; i++) {

			DriverObject->MajorFunction[i] = NtModelDrvCommon;
		}

		DriverObject->MajorFunction[IRP_MJ_CREATE]         = NtModelDrvCreateClose;
		DriverObject->MajorFunction[IRP_MJ_CLOSE]          = NtModelDrvCreateClose;
		DriverObject->MajorFunction[IRP_MJ_READ]           = NtModelDrvRead;
		DriverObject->MajorFunction[IRP_MJ_WRITE]          = NtModelDrvWrite;
		DriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL] = NtModelDrvIoControl; 

		KdPrintEx((77, 0, "Driver loaded successfully\n"));

		Status = STATUS_SUCCESS;

	} while (0);

	return Status;
}

VOID
DriverUnload(PDRIVER_OBJECT DriverObject)
{
	UNICODE_STRING SymbolicLinkName = { 0 };
	RtlInitUnicodeString(&SymbolicLinkName, SYMBOLIC_LINK_NAME);

	IoDeleteSymbolicLink(&SymbolicLinkName);

	IoDeleteDevice(DriverObject->DeviceObject);

	KdPrintEx((77, 0, "Driver uninstalled successfully\n"));
}

NTSTATUS
NtModelDrvCommon(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	Irp->IoStatus.Status      = STATUS_SUCCESS;
	Irp->IoStatus.Information = 0;

	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}

NTSTATUS
NtModelDrvCreateClose(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	Irp->IoStatus.Status      = STATUS_SUCCESS;
	Irp->IoStatus.Information = 0;

	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}

NTSTATUS
NtModelDrvRead(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	PIO_STACK_LOCATION Stack = IoGetCurrentIrpStackLocation(Irp);
	PVOID Buffer             = Irp->AssociatedIrp.SystemBuffer;
	ULONG_PTR BufferLen      = Stack->Parameters.Read.Length;

	TCHAR BeReadContent[] = L"Hi, r3, I come from r0";
	
	ULONG_PTR ReadLength = BufferLen > _tcslen(BeReadContent) * sizeof(TCHAR) ? _tcslen(BeReadContent) * sizeof(TCHAR) : BufferLen;

	RtlCopyMemory(Buffer, BeReadContent, ReadLength);

	Irp->IoStatus.Status      = STATUS_SUCCESS;
	Irp->IoStatus.Information = ReadLength;

	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}

NTSTATUS
NtModelDrvWrite(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	PIO_STACK_LOCATION Stack = IoGetCurrentIrpStackLocation(Irp);
	PVOID Buffer             = Irp->AssociatedIrp.SystemBuffer;
	ULONG_PTR BufferLen      = Stack->Parameters.Write.Length;

	WCHAR *WriteContent = ExAllocatePool(NonPagedPool, BufferLen);
	RtlCopyMemory(WriteContent, Buffer, BufferLen);

	UNICODE_STRING PrintString;
	RtlInitUnicodeString(&PrintString, WriteContent);

	KdPrintEx((77, 0, "%wZ\n", PrintString));

	ExFreePool(WriteContent);

	Irp->IoStatus.Status      = STATUS_SUCCESS;
	Irp->IoStatus.Information = BufferLen;

	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}

NTSTATUS
NtModelDrvIoControl(PDEVICE_OBJECT DeviceObject, PIRP Irp)
{
	PIO_STACK_LOCATION Stack = IoGetCurrentIrpStackLocation(Irp);


	switch (Stack->Parameters.DeviceIoControl.IoControlCode) {

	case IOCTL_NTMODELDRV_PRINTHELLOWORLD_STATS:
		KdPrintEx((77, 0, "Hello World\n"));
		break;

	default:
		KdPrintEx((77, 0, "Unknown iocontrol\n"));
		break;
	}
	
	Irp->IoStatus.Status      = STATUS_SUCCESS;
	Irp->IoStatus.Information = 0;

	IoCompleteRequest(Irp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}